import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.preprocessing import LabelEncoder, MinMaxScaler

# Define file paths
wpx_medal_counts_file = 'summerOly_medal_counts.csv'
wpx_athletes_file = 'summerOly_athletes.csv'
wpx_hosts_file = 'summerOly_hosts.csv'
wpx_programs_file = 'summerOly_programs.csv'

def read_csv_files():
    """Read CSV files and handle exceptions."""
    global wpx_medal_counts, wpx_athletes, wpx_hosts, wpx_programs

    try:
        wpx_medal_counts = pd.read_csv(wpx_medal_counts_file)
    except FileNotFoundError:
        print(f"Error: The file {wpx_medal_counts_file} was not found. Please check the file path!")
        exit()
    except UnicodeDecodeError:
        print(f"Error: The file {wpx_medal_counts_file} encoding is not supported by UTF-8. Please try another encoding format!")
        exit()

    try:
        wpx_athletes = pd.read_csv(wpx_athletes_file)
    except FileNotFoundError:
        print(f"Error: The file {wpx_athletes_file} was not found. Please check the file path!")
        exit()
    except UnicodeDecodeError:
        print(f"Error: The file {wpx_athletes_file} encoding is not supported by UTF-8. Please try another encoding format!")
        exit()

    try:
        wpx_hosts = pd.read_csv(wpx_hosts_file)
    except FileNotFoundError:
        print(f"Error: The file {wpx_hosts_file} was not found. Please check the file path!")
        exit()
    except UnicodeDecodeError:
        print(f"Error: The file {wpx_hosts_file} encoding is not supported by UTF-8. Please try another encoding format!")
        exit()

    wpx_encodings = ['utf-8', 'GBK', 'ISO-8859-1', 'latin1']  # Common encodings
    wpx_programs = None
    for wpx_encoding in wpx_encodings:
        try:
            wpx_programs = pd.read_csv(wpx_programs_file, encoding=wpx_encoding)
            break  # Exit the loop if the file is read successfully
        except UnicodeDecodeError:
            continue  # Try the next encoding
        except FileNotFoundError:
            print(f"Error: The file {wpx_programs_file} was not found. Please check the file path!")
            exit()

    if wpx_programs is None:
        print(f"Error: Unable to read the file {wpx_programs_file}. Please check the file encoding format!")
        exit()

def data_integration():
    """Integrate data from different sources and perform necessary calculations."""
    global wpx_data

    # Calculate the number of athletes and medal rate for each country
    wpx_athlete_count = wpx_athletes.groupby(['Year', 'Team'])['Name'].nunique().reset_index()
    wpx_athlete_count.columns = ['Year', 'Team', 'Athlete_Count']

    wpx_medal_count = wpx_athletes[wpx_athletes['Medal'] != 'No medal'].groupby(['Year', 'Team'])[
        'Name'].nunique().reset_index()
    wpx_medal_count.columns = ['Year', 'Team', 'Medal_Count']

    # Merge athlete count and medal rate
    wpx_athlete_medal = pd.merge(wpx_athlete_count, wpx_medal_count, on=['Year', 'Team'], how='left')
    wpx_athlete_medal['Medal_Rate'] = wpx_athlete_medal['Medal_Count'] / wpx_athlete_medal['Athlete_Count']

    # Merge medal counts and athlete data
    wpx_medal_counts.rename(columns={'NOC': 'Team'}, inplace=True)  # Rename NOC column to Team
    wpx_data = pd.merge(wpx_medal_counts, wpx_athlete_medal, on=['Year', 'Team'], how='left')

    # Merge host country information
    wpx_hosts['Host'] = wpx_hosts['Host'].str.strip()  # Remove whitespace from Host column
    wpx_medal_counts['Team'] = wpx_medal_counts['Team'].str.strip()  # Remove whitespace from Team column

    # Merge host data
    wpx_data = pd.merge(wpx_data, wpx_hosts, on='Year', how='left')
    wpx_data['Is_Host'] = wpx_data['Host'] == wpx_data['Team']  # Check if team is the host country

    # Convert programs data from wide to long format
    wpx_programs_long = wpx_programs.melt(id_vars=['Sport', 'Discipline', 'Code', 'Sports Governing Body'],
                                          var_name='Year', value_name='Included')
    wpx_programs_long = wpx_programs_long[wpx_programs_long['Included'] == 1]  # Keep only included sports

    # Clean Year column by removing non-numeric characters
    wpx_programs_long['Year'] = wpx_programs_long['Year'].str.replace(r'\D', '', regex=True)  # Remove non-numeric characters
    wpx_programs_long['Year'] = wpx_programs_long['Year'].astype(int)  # Convert Year to integer

    # Calculate the number of sports for each Olympic Games
    wpx_program_counts = wpx_programs_long.groupby('Year')['Sport'].nunique().reset_index()
    wpx_program_counts.columns = ['Year', 'Sport_Count']

    # Merge number of sports data
    wpx_data = pd.merge(wpx_data, wpx_program_counts, on='Year', how='left')

def data_cleaning():
    """Clean the integrated data by handling missing values, removing duplicates, and correcting outliers."""
    global wpx_data

    # Handle missing values by filling with mean values
    wpx_data['Athlete_Count'] = wpx_data['Athlete_Count'].fillna(wpx_data['Athlete_Count'].mean())
    wpx_data['Medal_Count'] = wpx_data['Medal_Count'].fillna(wpx_data['Medal_Count'].mean())
    wpx_data['Medal_Rate'] = wpx_data['Medal_Rate'].fillna(wpx_data['Medal_Rate'].mean())
    wpx_data['Sport_Count'] = wpx_data['Sport_Count'].fillna(wpx_data['Sport_Count'].mean())

    # Remove duplicate records
    wpx_data = wpx_data.drop_duplicates()

    # Handle outliers using a custom function
    def wpx_correct_outliers(series, window=3):
        Q1 = series.quantile(0.25)
        Q3 = series.quantile(0.75)
        IQR = Q3 - Q1
        lower = Q1 - 1.5 * IQR
        upper = Q3 + 1.5 * IQR
        wpx_outliers = (series < lower) | (series > upper)
        # Replace outliers with rolling window mean and ensure non-negative values
        series[wpx_outliers] = np.maximum(series.rolling(window, min_periods=1).mean().shift(1), 0)
        return series

    wpx_data['Gold'] = wpx_correct_outliers(wpx_data['Gold'])
    wpx_data['Athlete_Count'] = wpx_correct_outliers(wpx_data['Athlete_Count'])
    wpx_data['Medal_Rate'] = wpx_correct_outliers(wpx_data['Medal_Rate'])
    wpx_data['Sport_Count'] = wpx_correct_outliers(wpx_data['Sport_Count'])

def feature_engineering():
    """Create new features from the existing data."""
    global wpx_data

    wpx_data['Gold_Ratio'] = wpx_data['Gold'] / wpx_data['Total']  # Calculate gold medal ratio
    wpx_data['Is_Host'] = wpx_data['Is_Host'].astype(int)  # Convert Is_Host to integer

    # Encode categorical variables
    wpx_label_encoder = LabelEncoder()
    wpx_data['Team_encoded'] = wpx_label_encoder.fit_transform(wpx_data['Team'])

def normalization():
    """Normalize the data using MinMaxScaler."""
    global wpx_data

    wpx_scaler = MinMaxScaler()
    wpx_data[['Gold', 'Athlete_Count', 'Medal_Rate', 'Sport_Count']] = wpx_scaler.fit_transform(
        wpx_data[['Gold', 'Athlete_Count', 'Medal_Rate', 'Sport_Count']]
    )

def visualization():
    """Visualize the data using various plots."""
    global wpx_data

    # Ensure the picture folder exists
    wpx_picture_folder = "./picture"
    if not os.path.exists(wpx_picture_folder):
        os.makedirs(wpx_picture_folder)

    # 5.1 Distribution of Gold Medals
    plt.figure(figsize=(10, 6))
    sns.histplot(wpx_data['Gold'], bins=20, kde=True, color='red')
    plt.title('Distribution of Gold Medals')
    plt.xlabel('Gold Medals')
    plt.ylabel('Frequency')
    plt.grid(True)
    plt.savefig(f"{wpx_picture_folder}/1.1.1.png")
    plt.show()

    # 5.2 Gold Medals vs Athlete Count (Violin Plot)
    plt.figure(figsize=(10, 6))
    sns.violinplot(x='Is_Host', y='Gold', data=wpx_data, palette='bright')
    plt.title('Gold Medals vs Host Status')
    plt.xlabel('Is Host')
    plt.xticks(ticks=[0, 1], labels=['No', 'Yes'])
    plt.ylabel('Gold Medals')
    plt.grid(True)
    plt.savefig(f"{wpx_picture_folder}/1.1.2.png")
    plt.show()

    # 5.3 Number of Sports Over Time (Radar Chart)
    wpx_program_counts = wpx_data.groupby('Year')['Sport_Count'].mean().reset_index()
    categories = wpx_program_counts['Year'].tolist()
    values = wpx_program_counts['Sport_Count'].tolist()

    # Make the plot circular by appending the first value to the end
    categories = [*categories, categories[0]]
    values = [*values, values[0]]

    angles = [n / float(len(categories)) * 2 * np.pi for n in range(len(categories))]

    plt.figure(figsize=(10, 6))
    ax = plt.subplot(111, polar=True)
    plt.xticks(angles[:-1], categories[:-1], color='grey', size=8)
    ax.plot(angles, values, linewidth=1, linestyle='solid', color='green')
    ax.fill(angles, values, 'b', alpha=0.1)
    plt.title('Number of Sports Over Time', y=1.1)  # Adjust the title position
    ax.grid(True)
    plt.savefig(f"{wpx_picture_folder}/1.1.3.png")
    plt.show()

    # 5.4 Gold Ratio vs Medal Rate (Scatter Plot)
    plt.figure(figsize=(10, 6))
    sns.scatterplot(
        x='Medal_Rate',
        y='Gold_Ratio',
        data=wpx_data,
        hue='Is_Host',
        palette='bright',  # Use a brighter color palette
        alpha=0.8,  # Adjust point transparency
        s=100  # Adjust point size
    )
    plt.title('Gold Ratio vs Medal Rate (Scatter Plot)')
    plt.xlabel('Medal Rate')
    plt.ylabel('Gold Ratio')
    plt.legend(title='Is Host')
    plt.grid(True)
    plt.savefig(f"{wpx_picture_folder}/1.1.4.png")
    plt.show()

    # 5.5 Gold Ratio vs Medal Rate (KDE Plot)
    plt.figure(figsize=(10, 6))
    sns.kdeplot(
        x='Medal_Rate',
        y='Gold_Ratio',
        data=wpx_data,
        levels=5,
        color='red',
        alpha=0.5
    )
    plt.title('Gold Ratio vs Medal Rate (KDE Plot)')
    plt.xlabel('Medal Rate')
    plt.ylabel('Gold Ratio')
    plt.grid(True)
    plt.savefig(f"{wpx_picture_folder}/1.1.5.png")
    plt.show()

# Main execution
if __name__ == "__main__":
    # Step 1: Read CSV files
    read_csv_files()

    # Step 2: Data Integration
    data_integration()

    # Step 3: Data Cleaning
    data_cleaning()

    # Step 4: Feature Engineering
    feature_engineering()

    # Step 5: Normalization
    normalization()

    # Step 6: Visualization
    visualization()